#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
extern void col_2_im_custom_do(uint32_t blockDim, void *stream, uint8_t *col, uint8_t *img, int32_t kH, int32_t kW, int32_t stride, int32_t padding);
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void col_2_im_custom(GM_ADDR col, GM_ADDR img, int32_t kH, int32_t kW, int32_t stride, int32_t padding);
#endif

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 8;

    int32_t N = 1;
    int32_t C = 1;
    int32_t H = 8;
    int32_t W = 256;
    int32_t kernel_H = 2;
    int32_t kernel_W = 2;
    int32_t stride_val = 1;
    int32_t padding_val = 0;

    // 计算col矩阵的尺寸（输入）
    int32_t out_H = (H + 2 * padding_val - kernel_H) / stride_val + 1;
    int32_t out_W = (W + 2 * padding_val - kernel_W) / stride_val + 1;
    int32_t L = out_H * out_W;
    int32_t input_channels = C * kernel_H * kernel_W;  // col矩阵的通道数
    
    // col矩阵大小（输入）和图像大小（输出）
    size_t inputByteSize = N * input_channels * L * sizeof(float);   // col矩阵
    size_t outputByteSize = N * C * H * W * sizeof(float);           // 还原的图像

#ifdef ASCENDC_CPU_DEBUG
    uint8_t *col = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *img = (uint8_t *)AscendC::GmAlloc(outputByteSize);

    int32_t kH = 2;
    int32_t kW = 2;
    int32_t stride = 1;
    int32_t padding = 0;

    // 读取col矩阵作为输入
    ReadFile("./input/input_col.bin", inputByteSize, col, inputByteSize);

    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    ICPU_RUN_KF(col_2_im_custom, blockDim, col, img, kH, kW, stride, padding); // use this macro for cpu debug

    // 输出还原的图像
    WriteFile("./output/output_img.bin", img, outputByteSize);

    AscendC::GmFree((void *)col);
    AscendC::GmFree((void *)img);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    uint8_t *colHost, *imgHost;
    uint8_t *colDevice, *imgDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&colHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&imgHost), outputByteSize));
    CHECK_ACL(aclrtMalloc((void **)&colDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&imgDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    // 读取col矩阵作为输入
    ReadFile("./input/input_col.bin", inputByteSize, colHost, inputByteSize);

    CHECK_ACL(aclrtMemcpy(colDevice, inputByteSize, colHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    int32_t kH = 2;
    int32_t kW = 2;
    int32_t stride = 1;
    int32_t padding = 0;

    col_2_im_custom_do(blockDim, stream, colDevice, imgDevice, kH, kW, stride, padding);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(imgHost, outputByteSize, imgDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    
    // 输出还原的图像
    WriteFile("./output/output_img.bin", imgHost, outputByteSize);

    CHECK_ACL(aclrtFree(colDevice));
    CHECK_ACL(aclrtFree(imgDevice));
    CHECK_ACL(aclrtFreeHost(colHost));
    CHECK_ACL(aclrtFreeHost(imgHost));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}